import numpy as np
import pandas as pd
import time, os, psutil
import matplotlib.pyplot as plt


# 生成局部平稳(带宽型)转移矩阵 A


def generate_banded_A(N, bandwidth=5, rng=None):
    if rng is None:
        rng = np.random.default_rng()  # 推荐用新生成器
    A = np.zeros((N, N))
    for i in range(N):
        for j in range(max(0, i - bandwidth), min(N, i + bandwidth + 1)):
            A[i, j] = rng.random() + 0.1   # ✅ 这里改成 random()
    A /= A.sum(axis=1, keepdims=True)
    return A


# 标准 Viterbi（NumPy 向量化）


def viterbi_numpy(obs, A, B, pi):
    N, T = A.shape[0], len(obs)
    logA = np.log(A + 1e-12)
    logB = np.log(B + 1e-12)
    logPi = np.log(pi + 1e-12)

    delta = logPi + logB[:, obs[0]]
    for t in range(1, T):
        scores = delta[:, None] + logA
        delta = np.max(scores, axis=0) + logB[:, obs[t]]
    return np.max(delta)


# 分段局部 Viterbi（Segmented）


def viterbi_segmented_numpy(obs, A, B, pi, L=200, M=10):
    N, T = A.shape[0], len(obs)
    K = int(np.ceil(T / L))
    logA = np.log(A + 1e-12)
    logB = np.log(B + 1e-12)
    logPi = np.log(pi + 1e-12)

    # 第一段
    end0 = min(L, T)
    delta = logPi + logB[:, obs[0]]
    for t in range(1, end0):
        scores = delta[:, None] + logA
        delta = np.max(scores, axis=0) + logB[:, obs[t]]

    topM = np.argsort(delta)[-M:]
    prev_scores = delta[topM]

    # 后续分段
    for seg in range(1, K):
        start = seg * L
        end = min((seg + 1) * L, T)
        length = end - start

        first = np.max(prev_scores[:, None] + logA[topM, :], axis=0) + logB[:, obs[start]]
        delta_seg = first

        for t in range(1, length):
            scores = delta_seg[:, None] + logA
            delta_seg = np.max(scores, axis=0) + logB[:, obs[start + t]]

        topM = np.argsort(delta_seg)[-M:]
        prev_scores = delta_seg[topM]

    return np.max(prev_scores)


# 单组实验（返回 DataFrame）

def measure_performance(N=400, M_obs=6, T=20000, L=200, M=10, bandwidth=3, seed=42):
    rng = np.random.default_rng(seed)
    A = generate_banded_A(N, bandwidth, rng)
    B = rng.dirichlet(np.ones(M_obs), N)
    pi = rng.dirichlet(np.ones(N))
    obs = rng.integers(0, M_obs, size=T)

    # 标准 Viterbi
    t0 = time.perf_counter()
    logP_std = viterbi_numpy(obs, A, B, pi)
    t_std = time.perf_counter() - t0

    # 分段 Viterbi
    t0 = time.perf_counter()
    logP_seg = viterbi_segmented_numpy(obs, A, B, pi, L=L, M=M)
    t_seg = time.perf_counter() - t0

    diff = logP_seg - logP_std

    df = pd.DataFrame([
        {"N": N, "Method": "Standard", "Time(s)": t_std, "LogProb": logP_std, "Diff": 0},
        {"N": N, "Method": "Segmented", "Time(s)": t_seg, "LogProb": logP_seg, "Diff": diff},
    ])
    return df


# 批量运行实验 + 可视化


def batch_experiment():
    N_list = [200, 400, 800, 1200, 1600, 2000]
    all_results = []

    for N in N_list:
        print(f"\n▶ Running experiment: N={N}")
        df = measure_performance(N=N, M_obs=6, T=20000, L=200, M=10, bandwidth=3)
        all_results.append(df)

    results = pd.concat(all_results, ignore_index=True)
    results.to_csv("viterbi_batch_results.csv", index=False)
    print("\n✅ Results saved to viterbi_batch_results.csv")
    print(results)


    # 可视化 1：运行时间对比

    plt.figure(figsize=(8,5))
    for method, group in results.groupby("Method"):
        plt.plot(group["N"], group["Time(s)"], marker="o", label=method)
    plt.xlabel("状态数 N")
    plt.ylabel("运行时间 (s)")
    plt.title("Viterbi 算法运行时间对比（局部平稳模型）")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("runtime_comparison_banded.png", dpi=300)
    plt.show()


    # 可视化 2：路径概率差

    seg = results[results["Method"]=="Segmented"]
    plt.figure(figsize=(8,5))
    plt.plot(seg["N"], seg["Diff"], marker="x", color="purple")
    plt.xlabel("状态数 N")
    plt.ylabel("Δ log P")
    plt.title("最优路径概率差随模型规模变化")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("probability_diff_banded.png", dpi=300)
    plt.show()

    # 可视化 3：加速比 (Speedup = T_std / T_seg)

    std_times = results[results["Method"]=="Standard"][["N","Time(s)"]].set_index("N")
    seg_times = results[results["Method"]=="Segmented"][["N","Time(s)"]].set_index("N")
    speedup = (std_times["Time(s)"] / seg_times["Time(s)"]).reset_index()

    plt.figure(figsize=(8,5))
    plt.plot(speedup["N"], speedup["Time(s)"], marker="^", color="green")
    plt.xlabel("状态数 N")
    plt.ylabel("Speedup (T_std / T_seg)")
    plt.title("分段 Viterbi 算法加速比")
    plt.axhline(1.0, color="gray", linestyle="--", linewidth=1)
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("speedup_curve_banded.png", dpi=300)
    plt.show()

if __name__ == "__main__":
    batch_experiment()
